from captcha.image import ImageCaptcha #pip install captcha
import numpy as np
import matplotlib.pyplot as plt
from PIL import Image #pillow
#python 操作图片，opencv、pillow
import random

number = ['0', '1', '2', '3', '4', '5', '6', '7', '8', '9']
alphabet = ['a', 'b', 'c', 'd', 'e', 'f', 'g', 'h', 'i', 'j', 'k', 'l', 'm', 'n', 'o', 'p', 'q', 'r', 's', 't', 'u',
            'v', 'w', 'x', 'y', 'z']
ALPHABET = ['A', 'B', 'C', 'D', 'E', 'F', 'G', 'H', 'I', 'J', 'K', 'L', 'M', 'N', 'O', 'P', 'Q', 'R', 'S', 'T', 'U',
            'V', 'W', 'X', 'Y', 'Z']
MAX_CAPTCHA=4
IMAGE_HEIGHT = 60
IMAGE_WIDTH = 160
# 文本转向量
#char_set = number + alphabet + ALPHABET + ['_']  # 如果验证码长度小于4, '_'用来补齐


#char_set 从上面生成 随机生成四位
def random_captcha_text(char_set, captcha_size=MAX_CAPTCHA):#captcha_size是验证码有多少位
    captcha_text = []
    for i in range(captcha_size):
        c = random.choice(char_set)#随机的取一个值
        captcha_text.append(c)
    return captcha_text

#生成验证码
def gen_captcha_text_and_image(char_set,img_width,img_height):
    image = ImageCaptcha(width=img_width,height=img_height)

    captcha_text = random_captcha_text(char_set)
    captcha_text = ''.join(captcha_text)
    #把字符生成图像验证码
    captcha = image.generate(captcha_text)
    # image.write(captcha_text, captcha_text + '.jpg')

    # 把图片转换为数组形式
    captcha_image = Image.open(captcha)
    captcha_image = np.array(captcha_image)

    #得到label和图像
    return captcha_text, captcha_image

#转换为灰度图
def convert2gray(img):
    if len(img.shape) > 2:
        gray = np.mean(img, -1)
        # 上面的转法较快，正规转法如下
        # r, g, b = img[:,:,0], img[:,:,1], img[:,:,2]
        # gray = 0.2989 * r + 0.5870 * g + 0.1140 * b
        return gray
    else:
        return img
        
#根据字母获取ascii码，并根据ascii码转换下标
def char2pos(c):
    if c =='_':
        k = 62
        return k
    k = ord(c)-48
    if k > 9:
        k = ord(c) - 55
        if k > 35:
            k = ord(c) - 61
            if k > 61:
                raise ValueError('No Map')
    return k
    
#文本转换为向量
def text2vec(text):
    text_len = len(text)
    if text_len > MAX_CAPTCHA:
        raise ValueError('验证码最长4个字符')
    #生成MAX_CAPTCHA * CHAR_SET_LEN的都是0的矩阵
    vector = np.zeros(MAX_CAPTCHA * CHAR_SET_LEN)
    for i, c in enumerate(text):
        idx = i * CHAR_SET_LEN + int(char2pos(c))
        vector[idx] = 1
    return vector
# 向量转回文本
def vec2text(vec):
    text=[]
    for i, c in enumerate(vec):
        #char_at_pos = i #c/63
        char_idx = c % CHAR_SET_LEN
        if char_idx < 10:
            char_code = char_idx + ord('0')
        elif char_idx <36:
            char_code = char_idx - 10 + ord('A')
        elif char_idx < 62:
            char_code = char_idx-  36 + ord('a')
        elif char_idx == 62:
            char_code = ord('_')
        else:
            raise ValueError('error')
        text.append(chr(char_code))
    return ''.join(text)
# 生成图像及对应文本，并将大小不是(60, 160, 3)进行重新生成
def wrap_gen_captcha_text_and_image():
    while True:
        text, image = gen_captcha_text_and_image()
        if image.shape == (60, 160, 3):
            return text, image
            
# 生成一个训练batch数据
def get_next_batch(batch_size=128):
    batch_x = np.zeros([batch_size, IMAGE_HEIGHT * IMAGE_WIDTH])
    batch_y = np.zeros([batch_size, MAX_CAPTCHA * CHAR_SET_LEN])
    for i in range(batch_size):
        text, image = wrap_gen_captcha_text_and_image()
        image = convert2gray(image)
        batch_x[i, :] = image.flatten() / 255  # (image.flatten()-128)/128  mean为0
        batch_y[i, :] = text2vec(text)
    return batch_x, batch_y